#include <stdio.h>
#include <string.h>
#include <omp.h>
#include "lua.h"
#include "luaT.h"
#include "TH/TH.h"
#include "lauxlib.h"
#include "lualib.h"

#include "gsl/gsl_rng.h"
#include "ransampl.h"

static int roulette_sampling(lua_State* L) {
	THDoubleTensor* data = THDoubleTensor_newContiguous(luaT_toudata(L, -1, "torch.DoubleTensor"));
	double* probs = THDoubleTensor_data(data);
	int nrow = THDoubleTensor_size(data, 0);
	int ncol = THDoubleTensor_size(data, 1);
	
	THLongTensor* sample = THLongTensor_newWithSize1d(ncol);
	THDoubleTensor* rand = THDoubleTensor_newWithSize1d(ncol);
	THDoubleTensor_uniform(rand, 0, 1);
	
	int icol;
	#pragma omp parallel for shared(sample, rand, ncol, nrow, probs) private(icol)
	for (icol = 0; icol < ncol; icol++) {
		double r = THDoubleTensor_get1d(rand, icol);
		//printf("% f", r); 
		double acc = probs[icol];
		int irow;
		for (irow = 1; irow < nrow + 1; irow++) {
			if ( r <= acc ) {
				THLongTensor_set1d(sample, icol, irow);
				break;
			}
			acc += probs[irow*ncol+icol];
		}
	}
	//printf("\n");

	THDoubleTensor_free(rand);
	luaT_pushudata(L, sample, "torch.LongTensor");
	return 1;
}

// compute u = v * A
// u[i,:] = v[i,:] * A[i,:,:]
static THDoubleTensor* multi_vm_(THDoubleTensor* v, THDoubleTensor* M) {
	long n = THDoubleTensor_size(M, 0);
	long n1 = THDoubleTensor_size(M, 2);

	THDoubleTensor* u = THDoubleTensor_newWithSize2d(n, n1);
	
	int i;
	#pragma omp parallel for shared(M, v, u) private(i)
	for (i = 0; i < n; i++) {
		THDoubleTensor* M_ = THDoubleTensor_newTranspose(THDoubleTensor_newSelect(M, 0, i), 0, 1);
		THDoubleTensor* v_ = THDoubleTensor_newSelect(v, 0, i);
		THDoubleTensor* u_ = THDoubleTensor_newSelect(u, 0, i);
		THDoubleTensor_addmv(u_, 0, u_, 1, M_, v_);
		
		//THFree(M_);
		//THFree(u_);
		//THFree(v_);
	}

	return u;
}

static int multi_vm(lua_State* L) {
	THDoubleTensor* M = luaT_toudata(L, -1, "torch.DoubleTensor");
	THDoubleTensor* v = luaT_toudata(L, -2, "torch.DoubleTensor");
	THDoubleTensor* u = multi_vm_(v, M);

	luaT_pushudata(L, u, "torch.DoubleTensor");
	return 1;
}

// compute u = M * v
// u[i,:] = M[i,:,:] * v[i,:]
static THDoubleTensor* multi_mv_(THDoubleTensor* M, THDoubleTensor* v) {
	long n = THDoubleTensor_size(M, 0);
	long n1 = THDoubleTensor_size(M, 1);

	THDoubleTensor* u = THDoubleTensor_newWithSize2d(n, n1);
	
	int i = 0;
	THDoubleTensor** M_ = malloc(n*sizeof(THDoubleTensor*));
	THDoubleTensor** v_ = malloc(n*sizeof(THDoubleTensor*));
	THDoubleTensor** u_ = malloc(n*sizeof(THDoubleTensor*));
	
	for (i = 0; i < n; i++) {
		M_[i] = THDoubleTensor_newSelect(M, 0, i);
		v_[i] = THDoubleTensor_newSelect(v, 0, i);
		u_[i] = THDoubleTensor_newSelect(u, 0, i);
	}

	#pragma omp parallel for shared(M_, u_, v_) private(i)
	for (i = 0; i < n; i++) {
		THDoubleTensor_addmv(u_[i], 0, u_[i], 1, M_[i], v_[i]);	
	}

	for (i = 0; i < n; i++) {
		THDoubleTensor_free(M_[i]);
		THDoubleTensor_free(v_[i]);
		THDoubleTensor_free(u_[i]);
	}
	free(M_);
	free(v_);
	free(u_);

	return u;
}

static int multi_mv(lua_State* L) {
	THDoubleTensor* M = luaT_toudata(L, 1, "torch.DoubleTensor");
	THDoubleTensor* v = luaT_toudata(L, 2, "torch.DoubleTensor");
	THDoubleTensor* u = multi_mv_(M, v);

	luaT_pushudata(L, u, "torch.DoubleTensor");
	return 1;
}



// indepedent chain metropolis-hastings
// states[:,i] is chain of states drawn from a transition distribution
// ratios[:,i] is unnormalized_true_prob / unnormalized_transition_prob
static void indep_chain_metropolis_hastings_(THLongTensor* states, THDoubleTensor* ratios) {
	long chainlen = THLongTensor_size(states, 0);
	long nsamples = THLongTensor_size(states, 1);

	THDoubleTensor* r = THDoubleTensor_newWithSize2d(chainlen, nsamples);
	THDoubleTensor_uniform(r, 0, 1);

	int i;
	#pragma omp parallel for shared(states, ratios, r) private(i)
	for (i = 0; i < nsamples; i++) {
		int j;
		for (j = 1; j < chainlen; j++) {
			long cur_state = THLongTensor_get2d(states, j-1, i);
			double cur_ratio = THDoubleTensor_get2d(ratios, j-1, i);
			//long proposed_state = THLongTensor_get2d(states, j, i);
			double proposed_ratio = THDoubleTensor_get2d(ratios, j, i);

			double r_ = THDoubleTensor_get2d(r, j, i);
			double accept_rate = proposed_ratio / cur_ratio;
			if (accept_rate > 1) accept_rate = 1;

			if (r_ > accept_rate) {
				THLongTensor_set2d(states, j, i, cur_state);
				THDoubleTensor_set2d(ratios, j, i, cur_ratio);
			}
		} 
	}

	THDoubleTensor_free(r);
}

static int indep_chain_metropolis_hastings(lua_State* L) {
	THLongTensor* states = luaT_toudata(L, -2, "torch.LongTensor");
	THDoubleTensor* ratios = luaT_toudata(L, -1, "torch.DoubleTensor");
	indep_chain_metropolis_hastings_(states, ratios);
	return 0;
}

// alias sampling
// prob is a vector
static ransampl_ws* init_alias_sampling_(double* prob, int n) {
	ransampl_ws* ws = ransampl_alloc(n);
	ransampl_set(ws, prob);

	return ws;
}

static int init_alias_sampling(lua_State* L) {
	THDoubleTensor* prob = luaT_toudata(L, -1, "torch.DoubleTensor");
	double* prob__ = THDoubleTensor_data(prob);
	int n = THDoubleTensor_size(prob, 0);
	ransampl_ws* ws = init_alias_sampling_(prob__, n);
	
	THIntStorage* alias_ = THIntStorage_newWithData(ws->alias, ws->n);
	THDoubleStorage* prob_ = THDoubleStorage_newWithData(ws->prob, ws->n);
	
	luaT_pushudata(L, alias_, "torch.IntStorage");
	luaT_pushudata(L, prob_, "torch.DoubleStorage");

	free(ws);
	return 2;
}

static THLongTensor* gen_alias_sampling_(int* alias, double* prob, int n, int nsample) {
	ransampl_ws *ws = malloc( sizeof(ransampl_ws) );
	ws->n = n;
	ws->alias = alias;
	ws->prob = prob;

	gsl_rng_env_setup();
	gsl_rng* rng = gsl_rng_alloc( gsl_rng_default );

	THLongTensor* sample = THLongTensor_newWithSize1d(nsample);
	int i;
	#pragma omp parallel for private(i) shared(rng,ws,sample)
	for (i = 0; i < nsample; i++) {
		int x = ransampl_draw( ws, gsl_rng_uniform(rng), gsl_rng_uniform(rng));
		THLongTensor_set1d(sample, i, x + 1); // note that counting starts at 1 (in Lua)
	}
	
	gsl_rng_free(rng);
	free(ws);
	return sample;
}

static int gen_alias_sampling(lua_State* L) {
	THIntStorage* alias_ = luaT_toudata(L, -3, "torch.IntStorage");
	THDoubleStorage* prob_ = luaT_toudata(L, -2, "torch.DoubleStorage");
	int nsample = lua_tonumber(L, -1);
	int n = THIntStorage_size(alias_);

	THLongTensor* sample = gen_alias_sampling_(THIntStorage_data(alias_), 
											THDoubleStorage_data(prob_),
											n, nsample);

	luaT_pushudata(L, sample, "torch.LongTensor");

	return 1;
}

// compute outer-product  M = u * v
// M[i,:,:] = u[i,:] * v[i,:]
static THDoubleTensor* multi_oprod_(THDoubleTensor* u, THDoubleTensor* v) {
	long n = THDoubleTensor_size(u, 0);
	long n1 = THDoubleTensor_size(u, 1);
	long n2 = THDoubleTensor_size(v, 1);

	THDoubleTensor* M = THDoubleTensor_newWithSize3d(n, n1, n2);
	THDoubleTensor_zero(M);

	int i = 0;
	#pragma omp parallel for shared(M, v, u) private(i)
	for (i = 0; i < n; i++) {
		THDoubleTensor* M_ = THDoubleTensor_newSelect(M, 0, i);
		THDoubleTensor* v_ = THDoubleTensor_newSelect(v, 0, i);
		THDoubleTensor* u_ = THDoubleTensor_newSelect(u, 0, i);
		THDoubleTensor_addr(M_, 1, M_, 1, u_, v_);
		
		//THFree(M_);
		//THFree(u_);
		//THFree(v_);
	}

	return M;
}

static int multi_oprod(lua_State* L) {
	THDoubleTensor* u = luaT_toudata(L, -2, "torch.DoubleTensor");
	THDoubleTensor* v = luaT_toudata(L, -1, "torch.DoubleTensor");
	THDoubleTensor* M = multi_oprod_(u, v);

	luaT_pushudata(L, M, "torch.DoubleTensor");
	return 1;
}

int luaopen_cutils(lua_State *L) {
	lua_register(L, "roulette_sampling", roulette_sampling);
	lua_register(L, "multi_vm", multi_vm);
	lua_register(L, "multi_mv", multi_mv); 
	lua_register(L, "multi_oprod", multi_oprod);
	lua_register(L, "indep_chain_metropolis_hastings", indep_chain_metropolis_hastings);
	lua_register(L, "init_alias_sampling", init_alias_sampling);
	lua_register(L, "gen_alias_sampling", gen_alias_sampling);
	return 0;
}
